from gmssl import sm3, func
import gm_work_key_manager

def generate_salt(length=64):
    """生成指定长度的随机盐（十六进制字符串）"""
    # 生成32字节的盐值，即64个十六进制字符
    return func.random_hex(length)

key_len = 48

def get_hmac_key_byte():
    """
    获取用于HMAC的密钥字节

    1. 从gm_work_key_manager获取工作密钥的十六进制字符串。
    2. 将十六进制字符串以utf-8编码转换
    3. 使用SM3 KDF算法从工作密钥生成新的密钥，并将其转换为字节。

    Returns:
        bytes: 用于HMAC的密钥字节
    """

    work_key = gm_work_key_manager.get_work_key_hex()
    # 将十六进制字符串以utf-8编码转换
    work_key_bytes = work_key.encode('utf-8')
    new_key_hex = sm3.sm3_kdf(work_key_bytes, key_len)
    new_key_byte = bytes.fromhex(new_key_hex)
    return new_key_byte

def hmac_with_salt(password, salt):
    """
    使用HMAC-SM3算法计算密码和盐值的哈希

    Args:
        password (str): 输入的密码
        salt (str): 随机生成的盐值

    Returns:
        str: 计算得到的HMAC哈希值（十六进制字符串）
    """
    data = (password + salt).encode('utf-8')
    import hmac
    key = get_hmac_key_byte()
    hmac_result_hex = hmac.new(key, data, digestmod="sm3").hexdigest()
    return hmac_result_hex

def hmac_verify_password(input_password, salt, stored_hash):
    """
    验证输入的密码是否与存储的HMAC哈希匹配

    Args:
        input_password (str): 输入的密码
        salt (str): 随机生成的盐值
        stored_hash (str): 存储的HMAC哈希值

    Returns:
        bool: 如果输入密码与存储的HMAC哈希匹配则返回True，否则返回False
    """
    input_hash = hmac_with_salt(input_password, salt)
    return constant_time_compare(input_hash, stored_hash)

def hash_with_salt(password, salt):
    """
    使用盐值对密码进行SM3哈希处理

    Args:
        password (str): 输入的密码
        salt (str): 随机生成的盐值

    Returns:
        str: 计算得到的SM3哈希值（十六进制字符串）
    """
    # 将密码和盐拼接起来，并编码为字节
    data = (password + salt).encode('utf-8')
    # 使用SM3算法计算哈希
    hash_value = sm3.sm3_hash(func.bytes_to_list(data))
    # 返回十六进制字符串形式的哈希值
    return hash_value

def verify_password(input_password, salt, stored_hash):
    """
    验证输入的密码是否与存储的SM3哈希匹配

    Args:
        input_password (str): 输入的密码
        salt (str): 随机生成的盐值
        stored_hash (str): 存储的SM3哈希值

    Returns:
        bool: 如果输入密码与存储的SM3哈希匹配则返回True，否则返回False
    """
    # 使用相同的盐值对输入密码进行哈希处理
    input_hash = hash_with_salt(input_password, salt)
    # 比较新生成的哈希与存储的哈希
    # 使用常量时间比较防止计时攻击
    return constant_time_compare(input_hash, stored_hash)

def constant_time_compare(val1, val2):
    """
    进行常量时间比较，防止计时攻击

    Args:
        val1 (str): 第一个比较值
        val2 (str): 第二个比较值

    Returns:
        bool: 如果两个值相等则返回True，否则返回False
    """
    if len(val1) != len(val2):
        return False
    result = 0
    for x, y in zip(val1, val2):
        result |= ord(x) ^ ord(y)
    return result == 0

# 示例验证函数：验证HMAC功能
def test_hmac_functionality():
    password = "1"
    salt = generate_salt()
    print(f"Generated Salt: {salt}")

    # 计算HMAC哈希
    stored_hmac = hmac_with_salt(password, salt)
    print(f"Stored HMAC Hash: {stored_hmac}")

    # 验证正确的密码
    is_valid_correct = hmac_verify_password(password, salt, stored_hmac)
    print("Verification with correct password:", is_valid_correct)

    # 验证错误的密码
    wrong_password = "wrong_password"
    is_valid_wrong = hmac_verify_password(wrong_password, salt, stored_hmac)
    print("Verification with wrong password:", is_valid_wrong)

# 示例验证函数：验证普通SM3功能
def test_sm3_functionality():
    password = "1"
    salt = generate_salt()
    print(f"Generated Salt: {salt}")

    # 计算SM3哈希
    stored_hash = hash_with_salt(password, salt)
    print(f"Stored SM3 Hash: {stored_hash}")

    # 验证正确的密码
    is_valid_correct = verify_password(password, salt, stored_hash)
    print("Verification with correct password:", is_valid_correct)

    # 验证错误的密码
    wrong_password = "wrong_password"
    is_valid_wrong = verify_password(wrong_password, salt, stored_hash)
    print("Verification with wrong password:", is_valid_wrong)

if __name__ == "__main__":
    # 运行示例验证函数
    test_hmac_functionality()
    test_sm3_functionality()



